from flask import Flask, request, send_from_directory
import os
import subprocess
from werkzeug.utils import secure_filename

app = Flask(__name__)

UPLOAD_FOLDER = 'uploads'#存用户上传视频的文件夹
RESULT_FOLDER = 'results'#存检测后的视频的文件夹
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)#这两个确保文件夹存在
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULT_FOLDER'] = RESULT_FOLDER

ALLOWED_EXTENSIONS = {'mp4', 'webm'}# 允许的文件扩展名

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return "No file part", 400

    file = request.files['file']
    if file.filename == '':
        return "No selected file", 400

    if file and allowed_file(file.filename):
        # 安全地保存上传的文件
        filename = secure_filename(file.filename)
        input_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(input_path)

        # 定义输出文件路径
        output_filename = f"detected_{filename}"
        output_path = os.path.join(app.config['RESULT_FOLDER'], output_filename)

        # 调用YOLOv8检测函数
        from detect import run_yolov8_detection
        run_yolov8_detection(input_path, output_path)

        # 返回检测后的视频文件
        return send_from_directory(app.config['RESULT_FOLDER'], output_filename, as_attachment=True)

    return "Invalid file type", 400


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)